import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import numpy as np
import torch
import pyqtgraph.opengl as gl

from PySide6 import QtWidgets, QtCore
from configs.config import load_config
from data_process.radar_reader import RadarReader
from data_process.rtk_reader import RtkReader
from utils.coordinate_transformer import CoordinateTransformer

VOXEL_SIZE = (0.2, 0.2, 0.2)
X_RANGE = (-100, 100)
Y_RANGE = (0, 100)
Z_RANGE = (-0.5, 3.5)
T_WINDOW = 1

SPATIAL_RADIUS = (13, 13, 5)
SPATIAL_DECAY = lambda d: 1.0 / (1.0 + d)
# TEMPORAL_DECAY = lambda t: 1.0 / (1.0 + t)
# TEMPORAL_DECAY = lambda t: 1.0 * (t == 0)
TEMPORAL_DECAY = lambda t: 1.0

COMBINE_FN = lambda cohen, field: field

DEFAULT_FRAME = 755

device = 'cuda' if torch.cuda.is_available() else 'cpu'

def make_kernel():
    rx, ry, rz = SPATIAL_RADIUS
    t_len = T_WINDOW
    spatial_decay = SPATIAL_DECAY
    temporal_decay = TEMPORAL_DECAY
    kx = np.arange(-rx//2+1, rx//2+1)
    ky = np.arange(-ry//2+1, ry//2+1)
    kz = np.arange(-rz//2+1, rz//2+1)
    kt = np.arange(t_len)
    grid = np.stack(np.meshgrid(kx, ky, kz, kt, indexing='ij'), axis=-1)
    dist_space = np.abs(grid[...,0]) + np.abs(grid[...,1]) + np.abs(grid[...,2])
    dist_time = np.abs(grid[...,3] - (t_len - 1))
    kernel = spatial_decay(dist_space) * temporal_decay(dist_time)
    kernel[rx//2, ry//2, rz//2, t_len-1] = 1.0
    kernel = torch.from_numpy(kernel).float().to(device)
    offsets = torch.stack(torch.where(kernel != 0), dim=1)
    weights = kernel[kernel != 0]
    center = torch.tensor([rx//2, ry//2, rz//2, T_WINDOW-1], device=device)
    rel_offsets = offsets - center
    return rel_offsets, weights

KERNEL_OFFSETS, KERNEL_WEIGHTS = make_kernel()

def reconstruct_kernel_from_offsets(offsets, weights, kernel_shape):
    """
    offsets: (N, 4) tensor，偏移量（相对中心）
    weights: (N,) 1D tensor，系数
    kernel_shape: (kx, ky, kz, kt) 4D tuple
    返回：完整稀疏的kernel（torch.tensor），中心即(....//2, t_len-1)
    """
    kernel = torch.zeros(kernel_shape, dtype=weights.dtype, device=weights.device)
    center = torch.tensor([s//2 for s in kernel_shape[:3]] + [kernel_shape[3]-1], device=weights.device)
    abs_offsets = offsets + center
    for i in range(offsets.shape[0]):
        x, y, z, t = abs_offsets[i]
        kernel[x, y, z, t] = weights[i]
    return kernel

kernel_shape = (SPATIAL_RADIUS[0], SPATIAL_RADIUS[1], SPATIAL_RADIUS[2], T_WINDOW)
full_kernel = reconstruct_kernel_from_offsets(KERNEL_OFFSETS, KERNEL_WEIGHTS, kernel_shape)

print("完整稀疏卷积核张量（中心非零）:")
print(full_kernel.cpu().numpy())

def get_voxel_indices_torch(points, x_range, y_range, z_range, voxel_size):
    arr = torch.from_numpy(points).to(device)
    ix = torch.floor((arr[:, 0] - x_range[0]) / voxel_size[0]).long()
    iy = torch.floor((arr[:, 1] - y_range[0]) / voxel_size[1]).long()
    iz = torch.floor((arr[:, 2] - z_range[0]) / voxel_size[2]).long()
    return torch.stack([ix, iy, iz], dim=1)

def points_to_tensor_shape(x_range, y_range, z_range, voxel_size, t_window):
    sx = int((x_range[1] - x_range[0]) / voxel_size[0])
    sy = int((y_range[1] - y_range[0]) / voxel_size[1])
    sz = int((z_range[1] - z_range[0]) / voxel_size[2])
    return sx, sy, sz, t_window

def fast_sparse_field_torch(points_per_frame, cohen_per_frame, x_range, y_range, z_range, voxel_size, kernel_offsets, kernel_weights):
    sx, sy, sz, t_window = points_to_tensor_shape(x_range, y_range, z_range, voxel_size, len(points_per_frame))
    field = torch.zeros((sx, sy, sz, t_window), dtype=torch.float32, device=device)
    for t, (pts, coh) in enumerate(zip(points_per_frame, cohen_per_frame)):
        if pts.shape[0] == 0:
            continue
        pts_torch = torch.from_numpy(pts).float().to(device)
        cohen_torch = torch.from_numpy(coh).float().to(device)
        vxyz = get_voxel_indices_torch(pts, x_range, y_range, z_range, voxel_size)
        t_indices = torch.full((pts_torch.shape[0], 1), t, dtype=torch.long, device=device)
        base_indices = torch.cat([vxyz, t_indices], dim=1)
        stamp_indices = base_indices[:, None, :] + kernel_offsets[None, :, :]
        flat_indices = stamp_indices.reshape(-1, 4)
        flat_weights = (cohen_torch[:, None] * kernel_weights[None, :]).reshape(-1)
        valid_mask = (
            (flat_indices[:,0] >= 0) & (flat_indices[:,0] < sx) &
            (flat_indices[:,1] >= 0) & (flat_indices[:,1] < sy) &
            (flat_indices[:,2] >= 0) & (flat_indices[:,2] < sz) &
            (flat_indices[:,3] >= 0) & (flat_indices[:,3] < t_window)
        )
        flat_indices = flat_indices[valid_mask]
        flat_weights = flat_weights[valid_mask]
        linear_idx = (
            flat_indices[:,0] * sy * sz * t_window +
            flat_indices[:,1] * sz * t_window +
            flat_indices[:,2] * t_window +
            flat_indices[:,3]
        )
        field_flat = field.view(-1)
        field_flat.scatter_add_(0, linear_idx, flat_weights)
    return field.cpu().numpy()

def calc_field(points_per_frame, cohen_per_frame, voxel_params):
    x_range, y_range, z_range, voxel_size, t_len = voxel_params
    field = fast_sparse_field_torch(
        points_per_frame, cohen_per_frame, x_range, y_range, z_range, voxel_size, KERNEL_OFFSETS, KERNEL_WEIGHTS
    )
    return field[..., -1]

def apply_field_to_points(points, cohen, field, x_range, y_range, z_range, voxel_size, combine_fn=COMBINE_FN):
    ix, iy, iz = get_voxel_indices_torch(points, x_range, y_range, z_range, voxel_size).cpu().T
    sx, sy, sz = field.shape
    mask = (ix >= 0) & (ix < sx) & (iy >= 0) & (iy < sy) & (iz >= 0) & (iz < sz)
    result = np.zeros(points.shape[0], dtype=np.float32)
    field_vals = np.zeros(points.shape[0], dtype=np.float32)
    field_vals[mask] = field[ix[mask], iy[mask], iz[mask]]
    result[mask] = combine_fn(cohen[mask], field_vals[mask])
    return result


class MyGLViewWidget(gl.GLViewWidget):
    def __init__(self, parent=None, viewer=None):
        super().__init__(parent)
        self.viewer = viewer
    def mousePressEvent(self, event):
        if event.button() == QtCore.Qt.LeftButton and self.viewer is not None:
            pos = event.position() if hasattr(event, "position") else event.localPos()
            x, y = int(pos.x()), int(pos.y())
            self.viewer.select_point(x, y)
        super().mousePressEvent(event)


class SingleFramePointCloudViewer(QtWidgets.QWidget):
    def __init__(self, reader, rtk_reader, file_list, buffer_size=5):
        super().__init__()
        self.reader = reader
        self.rtk_reader = rtk_reader
        self.file_list = file_list
        self.buffer_size = buffer_size
        self.x_range = X_RANGE
        self.y_range = Y_RANGE
        self.z_range = Z_RANGE
        self.voxel_size = VOXEL_SIZE

        layout = QtWidgets.QVBoxLayout(self)
        self.view = MyGLViewWidget(viewer=self)
        self.view.setCameraPosition(distance=10)
        layout.addWidget(self.view, stretch=10)

        self.frame_slider = QtWidgets.QSlider(QtCore.Qt.Horizontal)
        self.frame_slider.setMinimum(0)
        self.frame_slider.setMaximum(len(self.file_list) - 1)
        self.frame_slider.setValue(DEFAULT_FRAME)
        layout.addWidget(QtWidgets.QLabel("帧号"))
        layout.addWidget(self.frame_slider)

        sliders_layout = QtWidgets.QHBoxLayout()
        self.s_cohen = self.make_slider("cohen分位数", sliders_layout, 0, 100, 0)
        self.s_cohen2 = self.make_slider("cohen2分位数", sliders_layout, 20, 100, 20)
        layout.addLayout(sliders_layout)

        self.scatter = gl.GLScatterPlotItem()
        self.view.addItem(self.scatter)

        self.info_label = QtWidgets.QLabel("点信息：")
        layout.addWidget(self.info_label)

        self.cur_points = []
        self.cur_cohen = self.cur_cohen2 = self.cur_xyz = self.cur_z = None

        self.frame_slider.valueChanged.connect(self.load_frame)
        self.s_cohen.valueChanged.connect(self.update_points)
        self.s_cohen2.valueChanged.connect(self.update_points)

        self.load_frame()

    def make_slider(self, name, layout, minv, maxv, val):
        box = QtWidgets.QVBoxLayout()
        label = QtWidgets.QLabel(f"{name}: {val}")
        slider = QtWidgets.QSlider(QtCore.Qt.Vertical)
        slider.setMinimum(minv)
        slider.setMaximum(maxv)
        slider.setValue(val)

        # --- 加减按钮 ---
        btn_plus = QtWidgets.QPushButton("+")
        btn_minus = QtWidgets.QPushButton("-")

        def set_value(delta):
            v = slider.value() + delta
            if minv <= v <= maxv:
                slider.setValue(v)

        btn_plus.clicked.connect(lambda: set_value(1))
        btn_minus.clicked.connect(lambda: set_value(-1))

        slider.valueChanged.connect(lambda v: label.setText(f"{name}: {v}"))

        btn_layout = QtWidgets.QHBoxLayout()
        btn_layout.addWidget(btn_minus)
        btn_layout.addWidget(btn_plus)

        box.addWidget(label)
        box.addWidget(slider)
        box.addLayout(btn_layout)
        layout.addLayout(box)
        return slider

    def load_frame(self):
        idx = self.frame_slider.value()
        indices = [max(0, idx-i) for i in range(self.buffer_size)][::-1]
        points_per_frame, cohen_per_frame = [], []

        frame_ids = [os.path.splitext(os.path.basename(self.file_list[i]))[0] for i in indices]
        points_dict = {}
        rtk_dict = {}

        folder_path = os.path.dirname(self.file_list[0])
        rtk_reader = self.rtk_reader

        for idx_f, i in enumerate(indices):
            ps = self.reader._read_single_file(self.file_list[i])
            frame_id = frame_ids[idx_f]
            points_dict[frame_id] = ps
            rtk_path = os.path.join(folder_path, frame_id + '.rtk')
            if os.path.exists(rtk_path):
                rtk_dict[frame_id] = rtk_reader._read_single_file(rtk_path)
            else:
                rtk_dict[frame_id] = None

        ref_rtk = rtk_dict[frame_ids[-1]]
        points_dict_corr = CoordinateTransformer.transform_points(points_dict, rtk_dict, ref_rtk)

        for frame_id in frame_ids:
            ps_corr = points_dict_corr.get(frame_id, [])
            if len(ps_corr) == 0:
                arr = np.zeros((0, 4), dtype=np.float32)
            else:
                arr = np.array([[p.x, p.y, p.z, p.cohen_factor] for p in ps_corr], dtype=np.float32)
            points_per_frame.append(arr[:, :3])
            cohen_per_frame.append(arr[:, 3])

        voxel_params = (self.x_range, self.y_range, self.z_range, self.voxel_size, self.buffer_size)
        field = calc_field(points_per_frame, cohen_per_frame, voxel_params)

        ps = points_per_frame[-1]
        cohen = cohen_per_frame[-1]
        cohen2 = apply_field_to_points(ps, cohen, field, self.x_range, self.y_range, self.z_range, self.voxel_size, combine_fn=COMBINE_FN)

        self.cur_points = ps
        self.cur_cohen = cohen
        self.cur_cohen2 = cohen2
        self.cur_xyz = ps
        self.cur_z = ps[:, 2]
        self.s_cohen.setValue(0)
        self.s_cohen2.setValue(0)
        self.update_points()

    def mousePressEvent(self, event):
        print("clicked")
        if event.button() == QtCore.Qt.LeftButton:
            # 兼容不同PySide版本
            if hasattr(event, "position"):
                pos = event.position()
            else:
                pos = event.localPos()
            x, y = int(pos.x()), int(pos.y())
            self.select_point(x, y)
        super().mousePressEvent(event)

    def select_point(self, x, y):
        if self.cur_xyz is None or len(self.cur_xyz) == 0:
            self.info_label.setText("点信息：无")
            return

        # 投影点到屏幕坐标
        # GLViewWidget里projectionMatrix和viewMatrix都是QMatrix4x4
        proj = self.view.projectionMatrix()
        modelview = self.view.viewMatrix()
        width = self.view.width()
        height = self.view.height()

        # 将点云变为齐次坐标
        pts = np.concatenate([self.cur_xyz, np.ones((self.cur_xyz.shape[0], 1))], axis=1)
        # 注意pyqtgraph的QMatrix4x4 .data()是列主序
        proj_mat = np.array(proj.data()).reshape(4,4).T
        model_mat = np.array(modelview.data()).reshape(4,4).T
        mvp = proj_mat @ model_mat
        pts_proj = pts @ mvp.T
        pts_proj /= pts_proj[:, 3:4]

        # 坐标映射到屏幕像素（注意y轴方向和OpenGL一致）
        screen_xy = np.empty((len(pts_proj), 2))
        screen_xy[:,0] = (pts_proj[:,0] + 1) * 0.5 * width
        screen_xy[:,1] = (1 - (pts_proj[:,1] + 1) * 0.5) * height

        mouse_xy = np.array([x, y])
        dists = np.linalg.norm(screen_xy - mouse_xy, axis=1)
        min_idx = np.argmin(dists)
        # 容忍30像素
        if dists[min_idx] < 30:
            pt = self.cur_xyz[min_idx]
            c1 = self.cur_cohen[min_idx]
            c2 = self.cur_cohen2[min_idx]
            self.info_label.setText(
                f"点 {min_idx}: (x={pt[0]:.2f}, y={pt[1]:.2f}, z={pt[2]:.2f}) cohen={c1:.3f} cohen2={c2:.3f}"
            )
        else:
            self.info_label.setText("点信息：未命中（可缩放/旋转再试）")


    def update_points(self):
        idx = self.frame_slider.value()
        file_path = self.file_list[idx]
        fname = os.path.basename(file_path)
        ms = int(fname.split('_')[-1][:-4])
        total_seconds = ms // 1000
        milliseconds = ms % 1000
        hours = total_seconds // 3600
        minutes = (total_seconds % 3600) // 60
        seconds = total_seconds % 60

        if self.cur_xyz is None or len(self.cur_xyz) == 0:
            self.scatter.setData(pos=np.zeros((1, 3)), size=6, color=(1, 0, 0, 0.1))
            self.setWindowTitle(
                f"帧 {idx} | {fname} | {hours:02}:{minutes:02}:{seconds:02}.{milliseconds:03} | 无点")
            return

        c_th = np.percentile(self.cur_cohen, self.s_cohen.value())
        c2_th = np.percentile(self.cur_cohen2, self.s_cohen2.value())
        mask = (self.cur_cohen >= c_th) & (self.cur_cohen2 >= c2_th)

        pts = self.cur_xyz[mask]
        z = self.cur_z[mask]
        if pts.shape[0] == 0:
            pts = np.zeros((1, 3), dtype=np.float32)
            z = np.zeros((1,), dtype=np.float32)

        z_norm = (z - z.min()) / (np.ptp(z) + 1e-6)
        color = np.empty((pts.shape[0], 4), dtype=float)
        color[:, 0] = z_norm
        color[:, 1] = 1 - z_norm
        color[:, 2] = 0
        color[:, 3] = 0.8

        self.scatter.setData(pos=pts, size=4, color=color)
        self.setWindowTitle(
            f"帧 {idx} | {fname} | {hours:02}:{minutes:02}:{seconds:02}.{milliseconds:03} | 点数: {pts.shape[0]}")

if __name__ == '__main__':
    cfg = load_config("./configs/data_process.py")
    reader = RadarReader(cfg)
    rtk_reader = RtkReader(cfg)
    data_root = list(cfg.get('data_root_list').values())[0]
    folder_path = os.path.join(data_root, 'SARRadar')
    files = reader._get_file_list(folder_path)
    files.sort(key=lambda x: int(x.split('_')[-1].split('.')[0]))

    app = QtWidgets.QApplication(sys.argv)
    win = SingleFramePointCloudViewer(reader, rtk_reader, files, buffer_size=T_WINDOW)
    win.show()
    sys.exit(app.exec())
